import os.path
import time
import pandas as pd
import numpy as np
import cython_optimized_functions

def rankdata(a, method='average', *, axis=None):
    # this rankdata refer to scipy.stats.rankdata (https://github.com/scipy/scipy/blob/v1.9.1/scipy/stats/_stats_py.py#L9047-L9153)
    if method not in ('average', 'min', 'max', 'dense', 'ordinal'):
        raise ValueError('unknown method "{0}"'.format(method))

    if axis is not None:
        a = np.asarray(a)
        if a.size == 0:
            np.core.multiarray.normalize_axis_index(axis, a.ndim)
            dt = np.float64 if method == 'average' else np.int_
            return np.empty(a.shape, dtype=dt)
        return np.apply_along_axis(rankdata, axis, a, method)

    arr = np.ravel(np.asarray(a))
    algo = 'mergesort' if method == 'ordinal' else 'quicksort'
    sorter = np.argsort(arr, kind=algo)

    inv = np.empty(sorter.size, dtype=np.intp)
    inv[sorter] = np.arange(sorter.size, dtype=np.intp)

    if method == 'ordinal':
        return inv + 1

    arr = arr[sorter]
    obs = np.r_[True, arr[1:] != arr[:-1]]
    dense = obs.cumsum()[inv]

    if method == 'dense':
        return dense

    # cumulative counts of each unique value
    count = np.r_[np.nonzero(obs)[0], len(obs)]

    if method == 'max':
        return count[dense]

    if method == 'min':
        return count[dense - 1] + 1

    # average method
    return .5 * (count[dense] + count[dense - 1] + 1)


def returns(x):
    a = (np.diff(x, axis = 0, append=np.nan)/x)
    return np.append([a[-1]],a[:-1],axis = 0)

def stddev(x, window=10):
    a_rolled = np.lib.stride_tricks.sliding_window_view(x, window,axis = 0)
    return np.append(np.full([window-1,np.size(x, 1)],np.nan),np.std(a_rolled, axis=-1, ddof=1),axis = 0)
    # return np.std(a_rolled, axis=-1)

def rank(x):
    return rankdata(x,method='min',axis=1)/np.size(x, 1)

def ts_argmax(x, window=10):
    a_rolled = np.lib.stride_tricks.sliding_window_view(x, window,axis = 0)
    return np.append(np.full([window-1,np.size(x, 1)],np.nan),np.argmax(a_rolled, axis=-1) + 1,axis = 0)

data_dir = os.path.join(os.path.expanduser("~"), 'Workspace/CraftProjects/play-madq/data')
data = pd.read_csv(os.path.join(data_dir, 'dataPerformance.csv'))
print("---data---")
print(data)
df = data.pivot(index='tradetime', columns='securityid', values='close')
print("---pivot---")
print(df)
t1 = time.time()
# 转换成 numpy.ndarray
close_prices = df.values.astype(np.float64)

rets = returns(close_prices)
np.putmask(close_prices, rets < 0, cython_optimized_functions.rolling_std_optimized(rets, 20))
result_ndarray = cython_optimized_functions.rankdata(cython_optimized_functions.ts_argmax(close_prices ** 2, 5)) - 0.5

result = pd.DataFrame(result_ndarray, columns=df.columns)
result.index = df.index
t4 = time.time()
print(f"---total time---: {t4 - t1}")
print("---result---")
print(result)

# ---total time---: 0.265078067779541
result.to_csv(os.path.join(data_dir, 'result_numpy_origin.csv'))
